# Copyright (c) 2024 AI绘画生成器
# 作者：zhongliuyang
# 版权所有，保留所有权利
# 本软件基于dashscope API开发
# 未经授权，禁止商业使用和分发

from http import HTTPStatus
from urllib.parse import urlparse, unquote
from pathlib import PurePosixPath
import requests
from dashscope import ImageSynthesis
import os
from flask import Flask, request, jsonify, send_file
import uuid
import time

app = Flask(__name__)

# 在文件开头添加风格常量
STYLES = {
    'auto': '默认',
    'photography': '摄影',
    'portrait': '人像写真',
    '3d cartoon': '3D卡通',
    'anime': '动画',
    'oil painting': '油画',
    'watercolor': '水彩',
    'sketch': '素描',
    'chinese painting': '中国画',
    'flat illustration': '扁平插画'
}

@app.route('/generate', methods=['POST'])
def generate_image():
    # 获取JSON请求数据
    data = request.get_json()
    print(data)
    if not data or 'prompt' not in data:
        return jsonify({'error': '缺少提示词参数'}), 400
    
    prompt = data['prompt']
    size = data.get('size', '256*256')  # 如果未指定size，默认使用256x256
    style = data.get('style', 'auto')  # 获取风格参数，默认为auto
    
    try:
        rsp = ImageSynthesis.call(
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            model=ImageSynthesis.Models.wanx_v1,
            prompt=prompt,
            n=1,
            style=f'<{style}>' if style != 'auto' else None,  # 根据选择添加风格标签
            size=size)
        print(rsp)
        if rsp.status_code == HTTPStatus.OK:
            # 生成唯一文件名
            timestamp = int(time.time())
            unique_id = str(uuid.uuid4())[:8]
            file_name = f"generated_{timestamp}_{unique_id}.png"
            
            # 保存图片
            for result in rsp.output.results:
                image_url = result.url
                with open(f"./images/{file_name}", 'wb+') as f:
                    f.write(requests.get(image_url).content)
                
                return jsonify({
                    'status': 'success',
                    'file_name': file_name,
                    'download_url': f'/images/{file_name}'
                })
            
            # 添加这个返回，以防for循环为空
            return jsonify({
                'status': 'error',
                'message': '没有生成任何图片'
            }), 400
            
        else:
            return jsonify({
                'status': 'error',
                'code': rsp.code,
                'message': rsp.message
            }), rsp.status_code
            
    except Exception as e:
        return jsonify({
            'status': 'error',
            'message': str(e)
        }), 500

@app.route('/images/<filename>')
def get_image(filename):
    try:
        return send_file(f"images/{filename}")
    except FileNotFoundError:
        return jsonify({'error': '图片不存在'}), 404

@app.route('/')
def index():
    return send_file('templates/index.html')

if __name__ == '__main__':
    # 确保images目录存在
    os.makedirs('images', exist_ok=True)
    app.run(debug=False, port=5000,host='0.0.0.0')